import numpy as np

# data sets
mails = [['goal', 'drink', 'defence', 'performance', 'field'],
        ['variance', 'drink', 'defence'],
        ['tutor', 'speed', 'defence', 'performance'],
        ['goal', 'speed', 'defence', 'field'],
        ['goal', 'drink', 'performance', 'field'],
        ['variance', 'speed', 'performance', 'field'],
        ['tutor', 'variance', 'performance'],
        ['goal', 'tutor', 'speed', 'performance', 'field'],
        ['tutor', 'variance', 'defence'],
        [],
        ['variance', 'drink', 'performance']]
labels = [0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1]

# 测试邮件
testMails0 = ['goal', 'speed', 'drink', 'defence', 'field']
testMails1 = ['tutor', 'variance', 'drink', 'performance']
testMails = [testMails0, testMails1]

# 计算先验概率（即类别概率）
class_p1 = sum(labels) / len(labels)
class_p0 = 1 - class_p1

def createDicts(mails):
    vocab = set([])
    for mail in mails:
        vocab = vocab | set(mail)
    vocab = list(vocab)
    return vocab

def vectorMat(vocab, mails):
    vector_mat = []
    for mail in mails:
        vector = np.zeros(len(vocab))
        for word in mail:
            if word in vocab:
                vector[vocab.index(word)] = 1
            else:
                print("the word: %s is not in my Vocabulary!" % word)
        vector_mat.append(vector)
    vector_mat = np.array(vector_mat)
    return vector_mat

def trainNB(vector_mat, labels):
    vocab_size = len(vector_mat[1])
    num0 = np.ones(vocab_size) # 拉普拉斯平滑, 避免概率为0
    num1 = np.ones(vocab_size) # 拉普拉斯平滑, 避免概率为0
    for i in range(len(labels)):
        if labels[i] == 0:
            num0 += vector_mat[i, :]
        else:
            num1 += vector_mat[i, :]
    vect_p0 = np.log(num0 / sum(num0))
    vect_p1 = np.log(num1 / sum(num1))
    print("num0, num0/sum(num0)/p0: ", num0, sum(num0))
    print("p0: ", num0/sum(num0))
    print("num1, num1/sum(num1)/p1: ", num1, sum(num1))
    print("p1: ", num1/sum(num1))
    return vect_p0, vect_p1

def classifyNB(test_mat, vect_p0, vect_p1):
    for mail_vect in test_mat:
        sample_p0 = sum(mail_vect * vect_p0) + np.log(class_p0)
        sample_p1 = sum(mail_vect * vect_p1) + np.log(class_p1)
        print('sample_p0:', 10 ** sample_p0)
        print('sample_p1:', 10 ** sample_p1)
        if sample_p1 > sample_p0:
            print('inform: ', 1)
        else:
            print('sports: ', 0)

if __name__ == '__main__':

    # 第一步：“去重”创建字典
    vocab = createDicts(mails)
    print("vocab: ", vocab)

    # 第二步：根据训练数据，创建邮件向量矩阵
    vector_mat = vectorMat(vocab, mails)
    print("vector_mat: ", vector_mat)

    # 第三步：在各类别中，统计字典的词频
    vect_p0, vect_p1 = trainNB(vector_mat, labels)
    print("vect_p0, vect_p1: ", vect_p0, vect_p1)

    # 第四步：测试邮件向量化，预测类别
    vector_mat = vectorMat(vocab, testMails)
    print("test_vector_mat:\n", vector_mat)
    classifyNB(vector_mat, vect_p0, vect_p1)
